#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import json
import time
import logging
from typing import Dict, List, Any, Optional

# 导入规则引擎
from rule_engine import RuleEngine

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger("智能体系统")


class Agent:
    """智能体基类"""
    def __init__(self, config: Dict[str, Any]):
        self.name = config["name"]
        self.description = config["description"]
        self.role = config["role"]
        self.capabilities = config["capabilities"]
        self.system_prompt = config["system_prompt"]
        self.logger = logging.getLogger(f"智能体.{self.name}")
        self.logger.info(f"初始化智能体: {self.name} (角色: {self.role})")
    
    def process(self, task: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, Any]:
        """处理任务"""
        self.logger.info(f"处理任务: {task.get('type', '未知')}")
        # 这里是模拟处理，实际应用中会调用对应的AI模型
        return {
            "agent": self.name,
            "result": f"{self.name} 处理了任务: {task.get('content', '')}",
            "timestamp": time.time(),
            "confidence": 0.9
        }


class MCP:
    """多功能连接器(Multi-Function Connector)基类"""
    def __init__(self, config: Dict[str, Any]):
        self.name = config["name"]
        self.description = config["description"]
        self.functions = {func["name"]: func for func in config["functions"]}
        self.logger = logging.getLogger(f"MCP.{self.name}")
        self.logger.info(f"初始化MCP: {self.name}")
    
    def call_function(self, function_name: str, **kwargs) -> Any:
        """调用MCP中的函数"""
        if function_name not in self.functions:
            raise ValueError(f"函数 {function_name} 在 {self.name} 中不存在")
        
        self.logger.info(f"调用函数: {function_name}")
        # 这里是函数调用的模拟实现
        if function_name == "process_user_input":
            return self._process_user_input(**kwargs)
        elif function_name == "format_response":
            return self._format_response(**kwargs)
        elif function_name == "create_task":
            return self._create_task(**kwargs)
        elif function_name == "assign_task":
            return self._assign_task(**kwargs)
        elif function_name == "track_task":
            return self._track_task(**kwargs)
        else:
            return {"status": "success", "message": f"函数 {function_name} 被调用"}
    
    # 用户交互MCP的函数实现
    def _process_user_input(self, user_input: str, context: Dict[str, Any]) -> Dict[str, Any]:
        return {
            "processed_text": user_input,
            "intent": "query" if "?" in user_input else "statement",
            "sentiment": "neutral",
            "timestamp": time.time()
        }
    
    def _format_response(self, agent_response: Dict[str, Any], context: Dict[str, Any]) -> str:
        return f"[{agent_response['agent']}] {agent_response['result']}"
    
    # 任务管理MCP的函数实现
    def _create_task(self, task_type: str, parameters: Dict[str, Any], priority: int) -> str:
        task_id = f"task_{int(time.time())}_{priority}"
        return task_id
    
    def _assign_task(self, task_id: str, available_agents: List[str]) -> Dict[str, Any]:
        # 简单的轮询分配策略
        if available_agents:
            assigned_agent = available_agents[0]
            return {
                "task_id": task_id,
                "assigned_to": assigned_agent,
                "status": "assigned",
                "timestamp": time.time()
            }
        return {"status": "error", "message": "无可用智能体"}
    
    def _track_task(self, task_id: str) -> Dict[str, Any]:
        return {
            "task_id": task_id,
            "status": "in_progress",
            "progress": 0.5,
            "timestamp": time.time()
        }


class ContextManager:
    """上下文管理器"""
    def __init__(self, config: Dict[str, Any]):
        self.max_history_length = config["max_history_length"]
        self.relevance_threshold = config["context_relevance_threshold"]
        self.memory_persistence = config["memory_persistence"]
        self.key_fields = config["key_context_fields"]
        self.context_data = {}
        self.logger = logging.getLogger("上下文管理器")
        
        # 初始化扩展的上下文字段
        self._init_extended_fields(config)
    
    def _init_extended_fields(self, config: Dict[str, Any]):
        """初始化扩展的上下文字段"""
        # 添加会话元数据字段
        self.session_metadata_fields = config.get("session_metadata_fields", [
            "session_start_time",
            "last_active_time",
            "interaction_count",
            "user_preferences"
        ])
        
        # 添加任务相关字段
        self.task_metadata_fields = config.get("task_metadata_fields", [
            "task_type",
            "task_priority",
            "task_dependencies",
            "task_deadline"
        ])
        
        # 添加环境信息字段
        self.environment_fields = config.get("environment_fields", [
            "system_time",
            "client_info",
            "locale",
            "timezone"
        ])
    
    def update_context(self, user_id: str, update_data: Dict[str, Any]):
        """更新上下文"""
        if user_id not in self.context_data:
            self.context_data[user_id] = {
                "history": [],
                "current_state": {},
                "session_metadata": {
                    "session_start_time": time.time(),
                    "interaction_count": 0
                }
            }
        
        # 更新会话元数据
        self.context_data[user_id]["session_metadata"]["last_active_time"] = time.time()
        self.context_data[user_id]["session_metadata"]["interaction_count"] += 1
        
        # 更新关键字段
        for field in self.key_fields:
            if field in update_data:
                self.context_data[user_id]["current_state"][field] = update_data[field]
        
        # 更新会话元数据字段
        for field in self.session_metadata_fields:
            if field in update_data:
                self.context_data[user_id]["session_metadata"][field] = update_data[field]
        
        # 更新历史记录
        self.context_data[user_id]["history"].append({
            "timestamp": time.time(),
            "data": update_data,
            "interaction_id": f"{user_id}_{len(self.context_data[user_id]['history'])}_time_{int(time.time())}"
        })
        
        # 保持历史记录长度
        if len(self.context_data[user_id]["history"]) > self.max_history_length:
            self.context_data[user_id]["history"] = self.context_data[user_id]["history"][-self.max_history_length:]
    
    def get_context(self, user_id: str) -> Dict[str, Any]:
        """获取上下文"""
        if user_id not in self.context_data:
            return {
                "history": [], 
                "current_state": {},
                "session_metadata": {}
            }
        return self.context_data[user_id]
    
    def get_relevant_context(self, user_id: str, query: str) -> Dict[str, Any]:
        """获取与查询相关的上下文"""
        full_context = self.get_context(user_id)
        
        # 在实际应用中，这里应该实现更复杂的相关性分析
        # 这里简化处理，返回完整上下文
        return full_context
    
    def clear_context(self, user_id: str):
        """清除用户上下文"""
        if user_id in self.context_data:
            del self.context_data[user_id]
            self.logger.info(f"已清除用户 {user_id} 的上下文")
    
    def save_context(self, file_path: str):
        """保存上下文到文件（用于持久化）"""
        if self.memory_persistence:
            try:
                with open(file_path, 'w', encoding='utf-8') as f:
                    json.dump(self.context_data, f, ensure_ascii=False, indent=2)
                self.logger.info(f"上下文已保存到 {file_path}")
            except Exception as e:
                self.logger.error(f"保存上下文失败: {e}")
    
    def load_context(self, file_path: str):
        """从文件加载上下文（用于恢复）"""
        if self.memory_persistence:
            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    self.context_data = json.load(f)
                self.logger.info(f"从 {file_path} 加载上下文成功")
            except Exception as e:
                self.logger.error(f"加载上下文失败: {e}")


class AgentSystem:
    """智能体系统核心"""
    def __init__(self, config_path: str):
        self.logger = logging.getLogger("智能体系统核心")
        self.logger.info("加载系统配置...")
        
        # 加载配置
        with open(config_path, 'r', encoding='utf-8') as f:
            config = json.load(f)
        
        # 初始化智能体
        self.agents = {}
        for agent_config in config["agents"]:
            agent = Agent(agent_config)
            self.agents[agent.name] = agent
        
        # 初始化MCP
        self.mcps = {}
        for mcp_config in config["mcps"]:
            mcp = MCP(mcp_config)
            self.mcps[mcp.name] = mcp
        
        # 初始化上下文管理器
        self.context_manager = ContextManager(config["context_config"])
        
        # 初始化规则引擎
        self.rules = config["rules"]
        self.rule_engine = RuleEngine(self.rules)
        
        # 添加系统级别的自定义规则
        self._add_system_rules()
        
        self.logger.info("系统初始化完成")
    
    def _add_system_rules(self):
        """添加系统级别的自定义规则"""
        # 添加任务优先级规则
        def task_priority_rule(context: Dict[str, Any]) -> bool:
            priority = context.get("priority", 1)
            return 1 <= priority <= 5
        
        self.rule_engine.add_custom_rule(
            "task_priority",
            "任务优先级必须在1-5之间",
            task_priority_rule,
            priority=85
        )
        
        # 添加上下文长度规则
        def context_length_rule(context: Dict[str, Any]) -> bool:
            user_input = context.get("text", "")
            return len(user_input) <= 1000  # 限制输入长度
        
        self.rule_engine.add_custom_rule(
            "context_length",
            "输入长度不能超过1000个字符",
            context_length_rule,
            priority=75
        )
    
    def handle_request(self, user_id: str, user_input: str) -> str:
        """处理用户请求"""
        self.logger.info(f"处理用户 {user_id} 的请求")
        
        # 创建规则检查上下文
        rule_context = {
            "text": user_input,
            "user_id": user_id,
            "start_time": time.time(),
            "request_source": "user_direct",
            "sensitive_fields": ["credit_card", "password", "ssn"]
        }
        
        # 检查是否应该阻止操作（安全检查）
        if self.rule_engine.should_block_action(rule_context):
            violations = self.rule_engine.get_violations()
            self.logger.warning(f"请求被安全规则阻止: {violations}")
            return "出于安全考虑，您的请求无法处理。"
        
        # 获取或创建用户上下文
        context = self.context_manager.get_context(user_id)
        
        # 使用用户交互MCP处理输入
        user_mcp = self.mcps["用户交互MCP"]
        processed_input = user_mcp.call_function("process_user_input", user_input=user_input, context=context)
        
        # 使用任务管理MCP创建任务
        task_mcp = self.mcps["任务管理MCP"]
        task_id = task_mcp.call_function(
            "create_task", 
            task_type=processed_input["intent"],
            parameters={"content": user_input},
            priority=1
        )
        
        # 根据任务类型分配给合适的智能体
        # 这里简化处理，实际应用中应该根据任务内容选择最合适的智能体
        assigned_result = task_mcp.call_function(
            "assign_task",
            task_id=task_id,
            available_agents=list(self.agents.keys())
        )
        
        # 让智能体处理任务
        target_agent_name = assigned_result.get("assigned_to")
        if target_agent_name and target_agent_name in self.agents:
            # 创建通信路径上下文（用于协调规则检查）
            rule_context["communication_path"] = ["用户交互MCP", "任务管理MCP", target_agent_name]
            
            # 检查所有规则
            if not self.rule_engine.check_rules(rule_context):
                violations = self.rule_engine.get_violations()
                self.logger.warning(f"规则检查失败: {violations}")
                # 即使规则检查失败，也继续处理，但记录警告
            
            agent_result = self.agents[target_agent_name].process(
                task={"type": processed_input["intent"], "content": user_input},
                context=context
            )
            
            # 检查输出准确性规则
            output_rule_context = {
                "confidence": agent_result.get("confidence", 1.0),
                "has_source": True,
                "requires_source": False
            }
            
            if not self.rule_engine.check_rules(output_rule_context):
                output_violations = self.rule_engine.get_violations()
                self.logger.warning(f"输出规则检查失败: {output_violations}")
            
            # 格式化响应
            response = user_mcp.call_function("format_response", agent_response=agent_result, context=context)
            
            # 更新上下文，包含更多元数据
            self.context_manager.update_context(user_id, {
                "user_identity": user_id,
                "conversation_topic": self._extract_topic(user_input),
                "task_progress": {"task_id": task_id, "status": "completed"},
                "recent_actions": [f"{target_agent_name} 处理了任务 {task_id}"],
                "last_response": response,
                "interaction_count": context.get("session_metadata", {}).get("interaction_count", 0) + 1,
                "user_preferences": {"response_format": "text"}
            })
            
            return response
        else:
            return "抱歉，无法处理您的请求，系统中没有可用的智能体。"
    
    def _extract_topic(self, text: str) -> str:
        """提取对话主题（简化实现）"""
        return text[:20] + "..." if len(text) > 20 else text
    
    def run_demo(self):
        """运行演示"""
        self.logger.info("开始系统演示")
        print("\n=== 智能体系统演示 ===")
        print("输入 'exit' 退出演示")
        print("输入 'status' 查看规则引擎状态")
        print("输入 'clear' 清除上下文")
        
        user_id = "demo_user"
        
        while True:
            try:
                user_input = input("\n您: ")
                
                if user_input.lower() == 'exit':
                    break
                elif user_input.lower() == 'status':
                    # 显示规则引擎状态
                    rule_status = self.rule_engine.get_rule_status()
                    print("\n--- 规则引擎状态 ---")
                    print(f"总规则数: {rule_status['total_rules']}")
                    print("规则详情:")
                    for rule in rule_status['rule_details']:
                        print(f"  - {rule['name']} (优先级: {rule['priority']})")
                    continue
                elif user_input.lower() == 'clear':
                    # 清除上下文
                    self.context_manager.clear_context(user_id)
                    print("上下文已清除")
                    continue
                
                response = self.handle_request(user_id, user_input)
                print(f"系统: {response}")
                
            except KeyboardInterrupt:
                break
        
        # 退出前保存上下文
        self.context_manager.save_context(f"context_backup_{user_id}.json")
        print("\n演示结束")


if __name__ == "__main__":
    try:
        # 创建系统实例
        system = AgentSystem("agents_config.json")
        
        # 运行演示
        system.run_demo()
        
    except Exception as e:
        logger.error(f"系统运行出错: {e}", exc_info=True)